# This file manages all gtest-based unit tests for triton kernels
# All tests are built here and can be run independently

# Get common variables from parent (if not already set)
if(NOT DEFINED ASCEND_HOME_PATH)
    set(ASCEND_HOME_PATH $ENV{ASCEND_HOME_PATH})
endif()

if(NOT DEFINED PLATFORM)
    if(CMAKE_SYSTEM_PROCESSOR STREQUAL "x86_64")
        set(PLATFORM "x86_64-linux")
    elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64" OR CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
        set(PLATFORM "aarch64-linux")
    else()
        message(WARNING "Unsupported architecture: ${CMAKE_SYSTEM_PROCESSOR}")
        set(PLATFORM "unknown")
    endif()
endif()

find_package(Python 3.11 REQUIRED COMPONENTS Interpreter Development)


# 关键修改：设置与 PyTorch 相同的 C++ ABI
# 检查 PyTorch 使用的 ABI 设置
# execute_process(
#     COMMAND python -c "import torch; print(torch._C._GLIBCXX_USE_CXX11_ABI)"
#     OUTPUT_VARIABLE TORCH_CXX11_ABI
#     OUTPUT_STRIP_TRAILING_WHITESPACE
# )

# message(STATUS "PyTorch uses CXX11 ABI: ${TORCH_CXX11_ABI}")

# 根据 PyTorch 的设置来配置 ABI
add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=0)
# message(STATUS "Using CXX11 ABI")


# Common include directories for all tests
set(TEST_INCLUDE_DIRS
    $<BUILD_INTERFACE:${ASCEND_HOME_PATH}/${PLATFORM}/include>
    $<BUILD_INTERFACE:${ASCEND_HOME_PATH}/${PLATFORM}/include/experiment/platform>
    $<BUILD_INTERFACE:${ASCEND_HOME_PATH}/${PLATFORM}/include/experiment/msprof>
    /usr/local/lib/python3.11/site-packages/torch/include
    /usr/local/lib/python3.11/site-packages/torch_npu/include
    /usr/local/lib/python3.11/site-packages/torch/include/torch/csrc/api/include
)

# Common link directories
set(TEST_LINK_DIRS
    ${ASCEND_HOME_PATH}/lib64
    /usr/local/lib/python3.11/site-packages/torch/lib
    /usr/local/lib/python3.11/site-packages/torch_npu/lib
)

# Get the binary directory path (relative to source root)
get_filename_component(BINARY_DIR "${CMAKE_SOURCE_DIR}/binary" ABSOLUTE)

# Test: npu_add_triton_test
add_executable(npu_add_triton_test
    npu_add_triton_test.cpp
)

target_include_directories(npu_add_triton_test PRIVATE
    ${TEST_INCLUDE_DIRS}
)

target_link_options(npu_add_triton_test PRIVATE
    "LINKER:--no-as-needed"
)

# Define binary directory path for test_utils
target_compile_definitions(npu_add_triton_test PRIVATE
    TEST_BINARY_DIR="${BINARY_DIR}"
)

# 设置 C++ 标准
set_target_properties(npu_add_triton_test PROPERTIES
    CXX_STANDARD 17
    CXX_STANDARD_REQUIRED ON
)

link_directories(${TEST_LINK_DIRS})

# 使用链接组解决循环依赖问题
target_link_libraries(npu_add_triton_test
    PRIVATE
        GTest::gtest
        GTest::gtest_main
        stdc++fs
)

# 将 PyTorch 相关库放在链接组中
target_link_options(npu_add_triton_test PRIVATE "LINKER:--start-group")
target_link_libraries(npu_add_triton_test
    PRIVATE
        "/usr/local/lib/python3.11/site-packages/torch/lib/libc10.so"
        torch
        torch_cpu
        torch_npu
        python3.11
        triton_kernel_loader
)
target_link_options(npu_add_triton_test PRIVATE "LINKER:--end-group")


# Enable testing
enable_testing()

# Add test to CTest
add_test(NAME NpuAddTritonTest COMMAND npu_add_triton_test)